#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""测试TypeInferencePass"""

from mcpy.compiler import MCPyCompiler
from mcpy.symbol_table import Symbol, TypeInfo, TypeKind
from typing import List, Optional


def type_matches(type_name: str, type_info: TypeInfo) -> bool:
    """检查类型名是否匹配 TypeInfo 对象

    Args:
        type_name: 类型名字符串，如 "int", "list[int]"
        type_info: TypeInfo 对象

    Returns:
        如果类型匹配返回 True，否则返回 False
    """
    return str(type_info) == type_name


class TestTypeInferencePass:
    """测试TypeInferencePass"""

    def _ensure_symbol_type_name(self, symbol: Optional[Symbol], type_name: str):
        """确保符号的类型"""

        self._ensure_symbol_type(
            symbol, lambda type_info: type_matches(type_name, type_info)
        )

    def _ensure_symbol_type(self, symbol: Optional[Symbol], cb):
        """确保符号的类型"""
        assert symbol is not None, f"符号应该存在"
        assert symbol.type_info is not None, f"符号 {symbol.name} 应该有类型信息"
        assert cb(symbol.type_info), f"符号 {symbol.name} 类型信息不满足条件"

    def _ensure_type_kind(self, type_info: Optional[TypeInfo], kind: TypeKind):
        """确保类型信息的类型"""
        assert type_info is not None, "类型信息应该存在"
        assert type_info.kind == kind, f"类型信息应该是 {kind}"

    def _ensure_type_args_name(
        self, type_info: Optional[TypeInfo], args_name: List[str]
    ):
        """确保类型信息的参数"""
        assert type_info is not None, "类型信息应该存在"
        assert [
            arg.name for arg in type_info.args
        ] == args_name, f"类型信息参数应该是 {args_name}"

    def _get_compiler(self):
        """获取编译器实例"""
        return MCPyCompiler()

    def test_basic_variable_assignment(self):
        """测试基本变量赋值的类型推断"""
        source = """
def test_func():
    x = 42
    y = "hello"
    z = 3.14
    return x, y, z
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查函数节点
        func_symbol = result.ir_module.lookup_symbol("test_func")
        assert func_symbol is not None, "应该找到函数符号"
        assert func_symbol.type_info is not None, "函数符号应该有类型信息"
        assert func_symbol.ir_node is not None, "函数符号应该有IR节点"

        # 检查函数局部变量类型
        x_symbol = func_symbol.ir_node.lookup_symbol("x")
        self._ensure_symbol_type_name(x_symbol, "int")

        y_symbol = func_symbol.ir_node.lookup_symbol("y")
        self._ensure_symbol_type_name(y_symbol, "str")

        z_symbol = func_symbol.ir_node.lookup_symbol("z")
        self._ensure_symbol_type_name(z_symbol, "float")

        # 检查函数返回值类型
        return_type = func_symbol.type_info.return_type
        self._ensure_type_kind(return_type, TypeKind.TUPLE)
        self._ensure_type_args_name(
            return_type,
            [
                "int",
                "str",
                "float",
            ],
        )

    def test_function_call_type_inference(self):
        """测试函数调用的类型推断"""
        source = """
def test_func():
    result1 = len([1, 2, 3])
    result2 = len("hello")
    result3 = abs(-5)
    result4 = abs(-3.14)
    return result1, result2, result3, result4
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查函数节点
        func_symbol = result.ir_module.lookup_symbol("test_func")
        assert func_symbol is not None, "应该找到函数符号"
        assert func_symbol.type_info is not None, "函数符号应该有类型信息"
        assert func_symbol.ir_node is not None, "函数符号应该有IR节点"

        # 检查函数局部变量类型
        result1_symbol = func_symbol.ir_node.lookup_symbol("result1")
        self._ensure_symbol_type_name(result1_symbol, "int")

        result2_symbol = func_symbol.ir_node.lookup_symbol("result2")
        self._ensure_symbol_type_name(result2_symbol, "int")

        result3_symbol = func_symbol.ir_node.lookup_symbol("result3")
        self._ensure_symbol_type_name(result3_symbol, "int")

        result4_symbol = func_symbol.ir_node.lookup_symbol("result4")
        self._ensure_symbol_type_name(result4_symbol, "float")

        # 检查函数返回值类型
        return_type = func_symbol.type_info.return_type
        self._ensure_type_kind(return_type, TypeKind.TUPLE)
        self._ensure_type_args_name(
            return_type,
            [
                "int",
                "int",
                "int",
                "float",
            ],
        )

    def test_class_method_type_inference(self):
        """测试类方法的类型推断"""
        source = """
class MyClass:
    def __init__(self, x: int):
        self.x = x

    def add(self, a: int, b: int):
        c = a + b
        return c + self.x + self.get_value()

    def get_value(self):
        return self.x
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查类符号
        class_symbol = result.ir_module.lookup_symbol("MyClass")
        assert class_symbol is not None, "应该找到类符号"
        assert class_symbol.type_info is not None, "类符号应该有类型信息"
        assert class_symbol.ir_node is not None, "类符号应该有IR节点"

        # 检查 MyClass.x 属性
        x_symbol = class_symbol.ir_node.lookup_symbol("x")
        self._ensure_symbol_type_name(x_symbol, "int")

        # 检查__init__方法
        init_symbol = class_symbol.ir_node.lookup_symbol("__init__")
        assert init_symbol is not None, "应该找到__init__方法符号"
        assert init_symbol.type_info is not None, "__init__方法符号应该有类型信息"
        assert init_symbol.ir_node is not None, "__init__方法符号应该有IR节点"

        # 检查add方法
        add_symbol = class_symbol.ir_node.lookup_symbol("add")
        assert add_symbol is not None, "应该找到add方法符号"
        assert add_symbol.type_info is not None, "add方法符号应该有类型信息"
        assert add_symbol.ir_node is not None, "add方法符号应该有IR节点"

        # 检查add方法的参数a
        a_symbol = add_symbol.ir_node.lookup_symbol("a")
        self._ensure_symbol_type_name(a_symbol, "int")

        # 检查add方法的参数b
        b_symbol = add_symbol.ir_node.lookup_symbol("b")
        self._ensure_symbol_type_name(b_symbol, "int")

        # 检查add方法的返回值类型
        add_return_type = add_symbol.type_info.return_type
        assert add_return_type is not None, "add方法应该有返回值类型"
        self._ensure_type_kind(add_return_type, TypeKind.PRIMITIVE)
        assert add_return_type.name == "int", "add方法返回值类型应该是int"

        # 检查get_value方法
        get_value_symbol = class_symbol.ir_node.lookup_symbol("get_value")
        assert get_value_symbol is not None, "应该找到get_value方法符号"
        assert get_value_symbol.type_info is not None, "get_value方法符号应该有类型信息"
        assert get_value_symbol.ir_node is not None, "get_value方法符号应该有IR节点"

        # 检查get_value方法的返回值类型
        get_value_return_type = get_value_symbol.type_info.return_type
        assert get_value_return_type is not None, "get_value方法应该有返回值类型"
        self._ensure_type_kind(get_value_return_type, TypeKind.PRIMITIVE)
        assert get_value_return_type.name == "int", "get_value方法返回值类型应该是int"

    def test_builtin_function_overload_resolution(self):
        """测试内建函数重载解析"""
        source = """
def test_overloads():
    # 测试len函数的重载
    list_len = len([1, 2, 3, 4])
    str_len = len("hello world")
    
    # 测试abs函数的重载
    int_abs = abs(-42)
    float_abs = abs(-3.14)
    
    return list_len, str_len, int_abs, float_abs
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查函数节点
        func_symbol = result.ir_module.lookup_symbol("test_overloads")
        assert func_symbol is not None, "应该找到函数符号"
        assert func_symbol.type_info is not None, "函数符号应该有类型信息"
        assert func_symbol.ir_node is not None, "函数符号应该有IR节点"

        # 检查函数局部变量类型 - len函数重载
        list_len_symbol = func_symbol.ir_node.lookup_symbol("list_len")
        self._ensure_symbol_type_name(list_len_symbol, "int")

        str_len_symbol = func_symbol.ir_node.lookup_symbol("str_len")
        self._ensure_symbol_type_name(str_len_symbol, "int")

        # 检查函数局部变量类型 - abs函数重载
        int_abs_symbol = func_symbol.ir_node.lookup_symbol("int_abs")
        self._ensure_symbol_type_name(int_abs_symbol, "int")

        float_abs_symbol = func_symbol.ir_node.lookup_symbol("float_abs")
        self._ensure_symbol_type_name(float_abs_symbol, "float")

        # 检查函数返回值类型
        return_type = func_symbol.type_info.return_type
        self._ensure_type_kind(return_type, TypeKind.TUPLE)
        self._ensure_type_args_name(
            return_type,
            [
                "int",
                "int",
                "int",
                "float",
            ],
        )

    def test_comprehensive_type_inference(self):
        """综合测试类型推断功能"""
        source = """
def comprehensive_test():
    # 基本类型
    num = 42
    text = "hello"
    flag = True
    
    # 函数调用
    length = len(text)
    absolute = abs(-10)
    
    # 二元操作
    sum_result = num + 10
    concat_result = text + " world"
    
    # 返回值
    return num, text, flag, length, absolute, sum_result, concat_result
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查函数符号
        func_symbol = result.ir_module.lookup_symbol("comprehensive_test")
        assert func_symbol is not None, "应该找到函数符号"
        assert func_symbol.type_info is not None, "函数符号应该有类型信息"
        assert func_symbol.ir_node is not None, "函数符号应该有IR节点"

        # 检查基本类型变量
        num_symbol = func_symbol.ir_node.lookup_symbol("num")
        self._ensure_symbol_type_name(num_symbol, "int")

        text_symbol = func_symbol.ir_node.lookup_symbol("text")
        self._ensure_symbol_type_name(text_symbol, "str")

        flag_symbol = func_symbol.ir_node.lookup_symbol("flag")
        self._ensure_symbol_type_name(flag_symbol, "bool")

        # 检查函数调用结果
        length_symbol = func_symbol.ir_node.lookup_symbol("length")
        self._ensure_symbol_type_name(length_symbol, "int")

        absolute_symbol = func_symbol.ir_node.lookup_symbol("absolute")
        self._ensure_symbol_type_name(absolute_symbol, "int")

        # 检查二元操作结果
        sum_result_symbol = func_symbol.ir_node.lookup_symbol("sum_result")
        self._ensure_symbol_type_name(sum_result_symbol, "int")

        concat_result_symbol = func_symbol.ir_node.lookup_symbol("concat_result")
        self._ensure_symbol_type_name(concat_result_symbol, "str")

        # 检查函数返回值类型
        return_type = func_symbol.type_info.return_type
        self._ensure_type_kind(return_type, TypeKind.TUPLE)
        self._ensure_type_args_name(
            return_type,
            [
                "int",
                "str",
                "bool",
                "int",
                "int",
                "int",
                "str",
            ],
        )

    def test_recursive_type_inference(self):
        """测试递归类型推断，这会引发编译错误"""
        source = """
class MyClass:
    def __init__(self, x: int):
        self.x = x

    def add(self, a: int, b: int):
        c = a + b
        return c + self.x + self.get_value()

    def get_value(self):
        # 递归调用 add 函数且不明确标注函数返回值类型
        # 这会触发类型检查错误，因为无法确定 add 函数的返回值类型
        return self.x + self.add(1, 2)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert not result.success
        assert any("循环引用检测到" in error for error in result.errors)

    def test_recursive_type_inference_with_annotation(self):
        """测试明确标注函数返回值类型避免递归类型推断"""
        source = """
class MyClass:
    def __init__(self, x: int):
        self.x = x

    def add(self, a: int, b: int):
        c = a + b
        return c + self.x + self.get_value()

    def get_value(self) -> int:
        # 递归调用 add 函数且明确标注函数返回值类型
        # 这不会触发类型检查错误，因为返回值类型已经标注为 int
        return self.x + self.add(1, 2)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success

    def test_binary_arithmetic_operations(self):
        """测试二元算术运算的类型推断"""
        source = """
def test_arithmetic():
    # 整数运算
    int_add = 5 + 3
    int_sub = 10 - 4
    int_mul = 6 * 7
    int_div = 8 / 2  # 除法总是返回float
    int_floordiv = 9 // 2
    int_mod = 10 % 3
    int_pow = 2 ** 3
    
    # 浮点数运算
    float_add = 5.5 + 2.3
    float_sub = 10.0 - 4.2
    float_mul = 3.14 * 2.0
    
    # 混合运算(int + float = float)
    mixed_add = 5 + 2.5
    mixed_mul = 3 * 1.5
    
    # 字符串连接
    str_concat = "hello" + " world"
    
    return (int_add, int_sub, int_mul, int_div, int_floordiv, int_mod, int_pow,
            float_add, float_sub, float_mul, mixed_add, mixed_mul, str_concat)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_arithmetic")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None
        assert func_symbol.type_info is not None

        # 检查整数运算结果类型
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("int_add"), "int"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("int_sub"), "int"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("int_mul"), "int"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("int_div"), "float"
        )  # 除法返回float
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("int_floordiv"), "int"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("int_mod"), "int"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("int_pow"), "int"
        )

        # 检查浮点数运算结果类型
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("float_add"), "float"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("float_sub"), "float"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("float_mul"), "float"
        )

        # 检查混合运算结果类型
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("mixed_add"), "float"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("mixed_mul"), "float"
        )

        # 检查字符串连接结果类型
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("str_concat"), "str"
        )

        self._ensure_type_kind(func_symbol.type_info.return_type, TypeKind.TUPLE)
        self._ensure_type_args_name(
            func_symbol.type_info.return_type,
            [
                "int",
                "int",
                "int",
                "float",
                "int",
                "int",
                "int",
                "float",
                "float",
                "float",
                "float",
                "float",
                "str",
            ],
        )

    def test_binary_comparison_operations(self):
        """测试二元比较运算的类型推断"""
        source = """
def test_comparison():
    # 数值比较
    eq_result = 5 == 5
    ne_result = 3 != 4
    lt_result = 2 < 3
    le_result = 4 <= 5
    gt_result = 6 > 5
    ge_result = 7 >= 7
    
    # 字符串比较
    str_eq = "hello" == "hello"
    str_ne = "foo" != "bar"
    
    return (eq_result, ne_result, lt_result, le_result, gt_result, ge_result, str_eq, str_ne)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_comparison")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 所有比较运算都应该返回bool类型
        comparison_vars = [
            "eq_result",
            "ne_result",
            "lt_result",
            "le_result",
            "gt_result",
            "ge_result",
            "str_eq",
            "str_ne",
        ]
        for var_name in comparison_vars:
            self._ensure_symbol_type_name(
                func_symbol.ir_node.lookup_symbol(var_name), "bool"
            )

    def test_binary_logical_operations(self):
        """测试二元逻辑运算的类型推断"""
        source = """
def test_logical():
    # 逻辑运算
    and_result = True and False
    or_result = True or False
    
    # 复合逻辑运算
    complex_and = (5 > 3) and (2 < 4)
    complex_or = (1 == 2) or (3 != 4)
    
    return (and_result, or_result, complex_and, complex_or)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_logical")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 所有逻辑运算都应该返回bool类型
        logical_vars = ["and_result", "or_result", "complex_and", "complex_or"]
        for var_name in logical_vars:
            self._ensure_symbol_type_name(
                func_symbol.ir_node.lookup_symbol(var_name), "bool"
            )

    def test_unary_operations(self):
        """测试一元运算的类型推断"""
        source = """
def test_unary():
    # 数值一元运算
    pos_int = +42
    neg_int = -42
    pos_float = +3.14
    neg_float = -3.14
    
    # 逻辑非运算
    not_true = not True
    not_false = not False
    not_expr = not (5 > 3)
    
    return (pos_int, neg_int, pos_float, neg_float, not_true, not_false, not_expr)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_unary")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查数值一元运算结果类型(保持原类型)
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("pos_int"), "int"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("neg_int"), "int"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("pos_float"), "float"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("neg_float"), "float"
        )

        # 检查逻辑非运算结果类型(总是bool)
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("not_true"), "bool"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("not_false"), "bool"
        )
        self._ensure_symbol_type_name(
            func_symbol.ir_node.lookup_symbol("not_expr"), "bool"
        )

    def test_collection_type_inference(self):
        """测试集合类型的类型推断"""
        source = """
def test_collections():
    # 列表类型推断
    int_list = [1, 2, 3]
    str_list = ["a", "b", "c"]
    mixed_list = [1, "hello", 3.14]  # 混合类型列表
    
    # 字典类型推断
    str_int_dict = {"a": 1, "b": 2}
    int_str_dict = {1: "one", 2: "two"}
    mixed_key_val_dict = {"a": 1, 2: "two", "c": 3.14}
    mixed_key_dict = {1: 1, "two": 2, "three": 3}
    mixed_val_dict = {1: 1, 2: True, 3: "three"}
    
    # 集合类型推断
    int_set = {1, 2, 3}
    str_set = {"x", "y", "z"}
    
    # 元组类型推断
    int_tuple = (1, 2, 3)
    mixed_tuple = (1, "hello", 3.14)
    
    return (int_list, str_list, mixed_list, str_int_dict, int_str_dict, 
            int_set, str_set, int_tuple, mixed_tuple)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_collections")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查列表类型
        int_list_symbol = func_symbol.ir_node.lookup_symbol("int_list")
        self._ensure_symbol_type_name(int_list_symbol, "list[int]")

        str_list_symbol = func_symbol.ir_node.lookup_symbol("str_list")
        self._ensure_symbol_type_name(str_list_symbol, "list[str]")

        mixed_list_symbol = func_symbol.ir_node.lookup_symbol("mixed_list")
        self._ensure_symbol_type_name(mixed_list_symbol, "list[Any]")

        # 检查字典类型
        str_int_dict_symbol = func_symbol.ir_node.lookup_symbol("str_int_dict")
        self._ensure_symbol_type_name(str_int_dict_symbol, "dict[str, int]")

        int_str_dict_symbol = func_symbol.ir_node.lookup_symbol("int_str_dict")
        self._ensure_symbol_type_name(int_str_dict_symbol, "dict[int, str]")

        mixed_key_val_dict_symbol = func_symbol.ir_node.lookup_symbol(
            "mixed_key_val_dict"
        )
        self._ensure_symbol_type_name(mixed_key_val_dict_symbol, "dict[Any, Any]")

        mixed_key_dict_symbol = func_symbol.ir_node.lookup_symbol("mixed_key_dict")
        self._ensure_symbol_type_name(mixed_key_dict_symbol, "dict[Any, int]")

        mixed_val_dict_symbol = func_symbol.ir_node.lookup_symbol("mixed_val_dict")
        self._ensure_symbol_type_name(mixed_val_dict_symbol, "dict[int, Any]")

        # 检查集合类型
        int_set_symbol = func_symbol.ir_node.lookup_symbol("int_set")
        self._ensure_symbol_type_name(int_set_symbol, "set[int]")

        str_set_symbol = func_symbol.ir_node.lookup_symbol("str_set")
        self._ensure_symbol_type_name(str_set_symbol, "set[str]")

        # 检查元组类型
        int_tuple_symbol = func_symbol.ir_node.lookup_symbol("int_tuple")
        self._ensure_symbol_type_name(int_tuple_symbol, "tuple[int, int, int]")

        mixed_tuple_symbol = func_symbol.ir_node.lookup_symbol("mixed_tuple")
        self._ensure_symbol_type_name(mixed_tuple_symbol, "tuple[int, str, float]")

    def test_method_call_type_inference(self):
        """测试方法调用的类型推断"""
        source = """
class Calculator:
    def __init__(self):
        self.value = 0
    
    def add(self, x: int) -> int:
        self.value += x
        return self.value
    
    def multiply(self, x: int) -> int:
        self.value *= x
        return self.value
    
    def get_value(self) -> int:
        return self.value

def test_method_calls():
    calc = Calculator()
    
    # 方法调用类型推断
    add_result = calc.add(5)
    multiply_result = calc.multiply(3)
    value_result = calc.get_value()
    
    return (add_result, multiply_result, value_result)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查类定义
        calc_class_symbol = result.ir_module.lookup_symbol("Calculator")
        assert calc_class_symbol is not None, "应该找到Calculator类符号"
        assert calc_class_symbol.type_info is not None, "Calculator类符号应该有类型信息"

        # 检查函数定义
        func_symbol = result.ir_module.lookup_symbol("test_method_calls")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None
        assert func_symbol.type_info is not None

        # 检查方法调用结果类型
        add_result_symbol = func_symbol.ir_node.lookup_symbol("add_result")
        self._ensure_symbol_type_name(add_result_symbol, "int")

        multiply_result_symbol = func_symbol.ir_node.lookup_symbol("multiply_result")
        self._ensure_symbol_type_name(multiply_result_symbol, "int")

        value_result_symbol = func_symbol.ir_node.lookup_symbol("value_result")
        self._ensure_symbol_type_name(value_result_symbol, "int")

        self._ensure_type_kind(func_symbol.type_info.return_type, TypeKind.TUPLE)
        self._ensure_type_args_name(
            func_symbol.type_info.return_type, ["int", "int", "int"]
        )

    def test_subscript_type_inference(self):
        """测试下标访问的类型推断"""
        source = """
def test_subscript():
    # 列表下标访问
    int_list = [1, 2, 3]
    list_element = int_list[0]
    
    # 字典下标访问
    str_dict = {"key1": "value1", "key2": "value2"}
    dict_value = str_dict["key1"]
    
    # 字符串下标访问
    text = "hello"
    char = text[0]
    
    return (list_element, dict_value, char)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_subscript")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查下标访问结果类型
        list_element_symbol = func_symbol.ir_node.lookup_symbol("list_element")
        self._ensure_symbol_type_name(list_element_symbol, "int")

        dict_value_symbol = func_symbol.ir_node.lookup_symbol("dict_value")
        self._ensure_symbol_type_name(dict_value_symbol, "str")

        char_symbol = func_symbol.ir_node.lookup_symbol("char")
        self._ensure_symbol_type_name(char_symbol, "str")

    def test_attribute_access_type_inference(self):
        """测试属性访问的类型推断"""
        source = """
class Person:
    def __init__(self, name: str, age: int):
        self.name = name
        self.age = age
        self.is_adult = age >= 18
    
    def get_info(self) -> str:
        return f"{self.name} is {self.age} years old"
    
    def get_age(self) -> int:
        return self.age

def test_attribute_access():
    person = Person("Alice", 25)
    
    # 基本属性访问类型推断
    person_name = person.name
    person_age = person.age
    adult_status = person.is_adult
    
    # 方法调用结果类型推断
    info_result = person.get_info()
    age_result = person.get_age()
    
    return (person_name, person_age, adult_status, info_result, age_result)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查类定义
        person_class_symbol = result.ir_module.lookup_symbol("Person")
        assert person_class_symbol is not None, "应该找到Person类符号"
        assert person_class_symbol.type_info is not None, "Person类符号应该有类型信息"

        # 检查函数定义
        func_symbol = result.ir_module.lookup_symbol("test_attribute_access")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查基本属性访问结果类型
        person_name_symbol = func_symbol.ir_node.lookup_symbol("person_name")
        self._ensure_symbol_type_name(person_name_symbol, "str")

        person_age_symbol = func_symbol.ir_node.lookup_symbol("person_age")
        self._ensure_symbol_type_name(person_age_symbol, "int")

        adult_status_symbol = func_symbol.ir_node.lookup_symbol("adult_status")
        self._ensure_symbol_type_name(adult_status_symbol, "bool")

        # 检查方法调用结果类型
        info_result_symbol = func_symbol.ir_node.lookup_symbol("info_result")
        self._ensure_symbol_type_name(info_result_symbol, "str")

        age_result_symbol = func_symbol.ir_node.lookup_symbol("age_result")
        self._ensure_symbol_type_name(age_result_symbol, "int")

    def test_nested_attribute_access_type_inference(self):
        """测试嵌套属性访问的类型推断"""
        source = """
class Person:
    def __init__(self, name: str, age: int):
        self.name = name
        self.age = age

class Address:
    def __init__(self, street: str, city: str):
        self.street = street
        self.city = city

class Employee:
    def __init__(self, person: Person, address: Address, salary: float):
        self.person = person
        self.address = address
        self.salary = salary

def test_nested_access():
    person = Person("Alice", 25)
    address = Address("123 Main St", "New York")
    employee = Employee(person, address, 50000.0)
    
    # 嵌套属性访问类型推断
    employee_name = employee.person.name
    employee_city = employee.address.city
    employee_salary = employee.salary
    
    return (employee_name, employee_city, employee_salary)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        # 检查函数定义
        func_symbol = result.ir_module.lookup_symbol("test_nested_access")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查直接属性访问结果类型
        employee_salary_symbol = func_symbol.ir_node.lookup_symbol("employee_salary")
        self._ensure_symbol_type_name(employee_salary_symbol, "float")

        # TODO: 嵌套属性访问暂时跳过，需要改进类型推断引擎
        # employee_name_symbol = func_symbol.ir_node.lookup_symbol("employee_name")
        # self._ensure_symbol_type_name(employee_name_symbol, "str")
        #
        # employee_city_symbol = func_symbol.ir_node.lookup_symbol("employee_city")
        # self._ensure_symbol_type_name(employee_city_symbol, "str")

    def test_comprehension_type_inference(self):
        """测试推导式的类型推断"""
        source = """
def test_comprehensions():
    # 列表推导式
    numbers = [1, 2, 3, 4, 5]
    squared = [x * x for x in numbers]
    squared1 = [i * x for i, x in enumerate(numbers)]

    strs = ["a", "b"]
    squared2 = [x for i, x in enumerate(strs)]
    
    # 字典推导式
    words = ["hello", "world"]
    word_lengths = {word: len(word) for word in words}
    
    # 集合推导式
    even_squares = {x * x for x in numbers if x % 2 == 0}
    
    return (squared, word_lengths, even_squares)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_comprehensions")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查推导式结果类型
        squared_symbol = func_symbol.ir_node.lookup_symbol("squared")
        self._ensure_symbol_type_name(squared_symbol, "list[int]")

        squared1_symbol = func_symbol.ir_node.lookup_symbol("squared1")
        self._ensure_symbol_type_name(squared1_symbol, "list[int]")

        squared2_symbol = func_symbol.ir_node.lookup_symbol("squared2")
        self._ensure_symbol_type_name(squared2_symbol, "list[str]")

        word_lengths_symbol = func_symbol.ir_node.lookup_symbol("word_lengths")
        self._ensure_symbol_type_name(word_lengths_symbol, "dict[str, int]")

        even_squares_symbol = func_symbol.ir_node.lookup_symbol("even_squares")
        self._ensure_symbol_type_name(even_squares_symbol, "set[int]")

    def test_generic_type_inference(self):
        """测试泛型类型推断"""
        source = """
from typing import List, Dict, Optional, Union

def test_generic_types():
    # 显式泛型类型注解
    int_list: List[int] = [1, 2, 3]
    str_dict: Dict[str, int] = {"a": 1, "b": 2}
    
    # Optional类型
    maybe_int: Optional[int] = 42
    maybe_none: Optional[str] = None
    
    # Union类型
    union_val: Union[int, str] = "hello"
    
    return (int_list, str_dict, maybe_int, maybe_none, union_val)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_generic_types")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查泛型类型
        int_list_symbol = func_symbol.ir_node.lookup_symbol("int_list")
        self._ensure_symbol_type_name(int_list_symbol, "list[int]")

        str_dict_symbol = func_symbol.ir_node.lookup_symbol("str_dict")
        self._ensure_symbol_type_name(str_dict_symbol, "dict[str, int]")

        # 检查Optional类型
        maybe_int_symbol = func_symbol.ir_node.lookup_symbol("maybe_int")
        self._ensure_symbol_type_name(maybe_int_symbol, "Optional[int]")

        maybe_none_symbol = func_symbol.ir_node.lookup_symbol("maybe_none")
        self._ensure_symbol_type_name(maybe_none_symbol, "Optional[str]")

        # 检查Union类型
        union_val_symbol = func_symbol.ir_node.lookup_symbol("union_val")
        self._ensure_symbol_type_name(union_val_symbol, "Union[int, str]")

    def test_nested_type_inference(self):
        """测试嵌套类型推断"""
        source = """
def test_nested_types():
    # 嵌套列表
    nested_list = [[1, 2], [3, 4], [5, 6]]
    
    # 嵌套字典
    nested_dict = {
        "group1": {"a": 1, "b": 2},
        "group2": {"c": 3, "d": 4}
    }
    
    # 复杂嵌套结构
    complex_data = {
        "users": [
            {"name": "Alice", "age": 25},
            {"name": "Bob", "age": 30}
        ],
        "settings": {"theme": "dark", "notifications": True}
    }
    user = complex_data.users[0]
    count = len(complex_data.users[0])
    
    return (nested_list, nested_dict, complex_data)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_nested_types")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查嵌套列表类型
        nested_list_symbol = func_symbol.ir_node.lookup_symbol("nested_list")
        self._ensure_symbol_type_name(nested_list_symbol, "list[list[int]]")

        # 检查嵌套字典类型
        nested_dict_symbol = func_symbol.ir_node.lookup_symbol("nested_dict")
        self._ensure_symbol_type_name(nested_dict_symbol, "dict[str, dict[str, int]]")

        # 检查复杂数据结构类型
        complex_data_symbol = func_symbol.ir_node.lookup_symbol("complex_data")
        self._ensure_symbol_type_name(complex_data_symbol, "dict[str, Any]")

        user = func_symbol.ir_node.lookup_symbol("user")
        self._ensure_symbol_type_name(user, "dict[str, Any]")

        count = func_symbol.ir_node.lookup_symbol("count")
        self._ensure_symbol_type_name(count, "int")

    def test_control_flow_type_inference(self):
        """测试控制流语句中的类型推断"""
        source = """
def test_control_flow():
    # if-else语句中的类型推断
    condition = True
    result = 0  # 先定义变量
    if condition:
        result = 42
    else:
        result = "hello"
    
    # for循环中的类型推断
    numbers = [1, 2, 3, 4, 5]
    total = 0
    for num in numbers:
        total += num
    
    # while循环中的类型推断
    counter = 0
    while counter < 10:
        counter += 1
    
    # 列表推导式中的类型推断
    squared_numbers = [x * x for x in numbers]
    
    return (result, total, counter, squared_numbers)
"""
        compiler = self._get_compiler()
        result = compiler.compile_source(source, "test.py")
        assert result.success
        assert result.ir_module is not None

        func_symbol = result.ir_module.lookup_symbol("test_control_flow")
        assert func_symbol is not None
        assert func_symbol.ir_node is not None

        # 检查条件变量类型
        condition_symbol = func_symbol.ir_node.lookup_symbol("condition")
        self._ensure_symbol_type_name(condition_symbol, "bool")

        # 检查循环变量类型
        total_symbol = func_symbol.ir_node.lookup_symbol("total")
        self._ensure_symbol_type_name(total_symbol, "int")

        counter_symbol = func_symbol.ir_node.lookup_symbol("counter")
        self._ensure_symbol_type_name(counter_symbol, "int")

        # 检查推导式结果类型
        squared_numbers_symbol = func_symbol.ir_node.lookup_symbol("squared_numbers")
        self._ensure_symbol_type_name(squared_numbers_symbol, "list[int]")


#     def test_exception_handling_type_inference(self):
#         """测试异常处理中的类型推断"""
#         source = """
# def test_exception_handling():
#     try:
#         value = int("42")
#         result = value * 2
#     except ValueError as e:
#         error_msg = str(e)
#         result = 0
#     except Exception as e:
#         general_error = str(e)
#         result = -1
#     finally:
#         cleanup_done = True

#     return (result, cleanup_done)
# """
#         compiler = self._get_compiler()
#         result = compiler.compile_source(source, "test.py")
#         assert result.success
#         assert result.ir_module is not None

#         func_symbol = result.ir_module.lookup_symbol("test_exception_handling")
#         assert func_symbol is not None
#         assert func_symbol.ir_node is not None

#         # 检查异常处理中的变量类型
#         value_symbol = func_symbol.ir_node.lookup_symbol("value")
#         self._ensure_symbol_type_name(value_symbol, "int")

#         cleanup_done_symbol = func_symbol.ir_node.lookup_symbol("cleanup_done")
#         self._ensure_symbol_type_name(cleanup_done_symbol, "bool")
